import openai
import json
from typing import List, Dict, Any, Optional

class BasicOpenAIAgent:
    """
    基础的OpenAI Agent类
    """

    def __init__(self, api_key: str, model: str = "gpt-3.5-turbo"):
        """
        初始化Agent

        Args:
            api_key (str): OpenAI API密钥
            model (str): 使用的模型，默认为gpt-3.5-turbo
        """
        self.api_key = api_key
        self.model = model
        self.client = openai.OpenAI(api_key=self.api_key)
        self.conversation_history: List[Dict[str, str]] = []

    def set_system_prompt(self, prompt: str):
        """
        设置系统提示词

        Args:
            prompt (str): 系统提示词
        """
        system_message = {"role": "system", "content": prompt}
        # 检查是否已有系统消息，如果有则替换，否则添加到开头
        if self.conversation_history and self.conversation_history[0]["role"] == "system":
            self.conversation_history[0] = system_message
        else:
            self.conversation_history.insert(0, system_message)

    def add_user_message(self, message: str):
        """
        添加用户消息到对话历史

        Args:
            message (str): 用户消息
        """
        self.conversation_history.append({"role": "user", "content": message})

    def add_assistant_message(self, message: str):
        """
        添加助手消息到对话历史

        Args:
            message (str): 助手消息
        """
        self.conversation_history.append({"role": "assistant", "content": message})

    def get_response(self, user_input: str, max_tokens: Optional[int] = None) -> str:
        """
        获取AI响应

        Args:
            user_input (str): 用户输入
            max_tokens (int, optional): 最大token数

        Returns:
            str: AI的响应
        """
        # 添加用户消息
        self.add_user_message(user_input)

        # 构造请求参数
        params = {
            "model": self.model,
            "messages": self.conversation_history
        }

        if max_tokens:
            params["max_tokens"] = max_tokens

        try:
            # 发送请求到OpenAI API
            response = self.client.chat.completions.create(**params)

            # 提取响应内容
            assistant_reply = response.choices[0].message.content

            # 添加助手回复到历史记录
            self.add_assistant_message(assistant_reply)

            return assistant_reply

        except Exception as e:
            return f"Error: {str(e)}"

    def reset_conversation(self):
        """
        重置对话历史（保留系统提示词）
        """
        system_messages = [msg for msg in self.conversation_history if msg["role"] == "system"]
        self.conversation_history = system_messages

    def get_conversation_history(self) -> List[Dict[str, str]]:
        """
        获取对话历史

        Returns:
            List[Dict[str, str]]: 对话历史
        """
        return self.conversation_history.copy()


# 使用示例
if __name__ == "__main__":
    # 初始化Agent（请替换为您的实际API密钥）

    agent = BasicOpenAIAgent(api_key="your-openai-api-key-here")

    # 设置系统提示词
    agent.set_system_prompt("你是一个乐于助人的AI助手，会用中文回答问题。")

    # 进行对话
    response = agent.get_response("你好，你能帮我做什么？")
    print(f"AI: {response}")

    # 继续对话
    response = agent.get_response("请告诉我一些关于人工智能的趣事。")
    print(f"AI: {response}")

    # 查看对话历史
    print("\n对话历史:")
    for msg in agent.get_conversation_history():
        print(f"{msg['role']}: {msg['content']}")
